import re
from until import datapath


def normal_str(s):
    """

    :param s: 原始字符串【'   You are5555 my best friend.  '】
    :return: 格式化之后的字符串【you are my best friend .】
    """
    s = s.lower().strip()  # 小写加去空格
    s = re.sub(r'([.!?])', r' \1', s)  # 匹配.!?，在其前面插入一个空格
    s = re.sub(r'[^a-zA-Z.!?]+', r' ', s) # 匹配非英文字符，将其替换为空格
    return s


# 构建英法语词典
def get_data():
    with open(datapath, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    # l = []
    # print(lines[0].strip().split('\t'))
    # for s in lines[0].strip().split('\t'):
    #     l.append(normal_str(s))
    #     print(l)
    my_pairs = []
    # lines[0]:    i m .	j ai ans .
    for line in lines:
        temp = []
        # line.strip().split('\t') : ['i m .', 'j ai ans .']
        for word in line.strip().split('\t'):
            # word: i m .  j ai ans .
            temp.append(normal_str(word))
            # temp: ['i m .', 'j ai ans .']
        my_pairs.append(temp)
        # my_pairs: [['i m .', 'j ai ans .'], ['i m .', 'j ai ans .']]

    english_word2index = {'SOS': 0, 'EOS': 1}
    french_word2index = {'SOS': 0, 'EOS': 1}

    for pair in my_pairs:
        # 'i m ok .'  ->  ['i', 'm', 'ok', '.']
        for word in pair[0].split(' '):
            if word not in english_word2index.keys():
                english_word2index[word] = len(english_word2index)
        for word in pair[1].split(' '):
            if word not in french_word2index.keys():
                french_word2index[word] = len(french_word2index)

    english_index2word = {v: k for k, v in english_word2index.items()}
    french_index2word = {v: k for k, v in french_word2index.items()}

    return english_word2index, english_index2word, len(english_word2index), \
        french_word2index, french_index2word, len(french_word2index), my_pairs


if __name__ == '__main__':
    # 测试normal_str()
    s = '   You are5555 my best friend!  '
    r = normal_str(s)
    print('格式化之后的字符串', r)

    # 测试get_data()
    english_word2index, english_index2word, english_word_n, \
        french_word2index, french_index2word, french_word_n, my_pairs = get_data()
    # print('英文词典：', english_word2index)
    # print('法文词典：', french_word2index)
    # print('格式化后的语料对：', my_pairs)
    print('英文词典大小：', english_word_n)
    print('法文词典大小：', french_word_n)
